#!/usr/bin/python
# -*- coding:utf-8 -*-

"""
更新任务漏洞等级
暂时不考虑更新这部分数据
"""

from app.db.models.tasks import Task


def update_task_level(task_id=None):
    level_map = {2: "低危", 3: "中危", 4: "高危", 5: "严重"}
    query = {"taskType__in": ["vul", "ssl", "securityEvent", "content"]}
    if task_id:
        query.update({"id__gt": task_id})
    all_task = Task.objects.filter(**query).timeout(False)
    for task in all_task:
        warnings = task.result.warnings
        if not warnings:
            continue
        is_change = False
        for warning in task.result.warnings:
            if (warning.level in level_map) and warning.levelName == level_map[warning.level]:
                break
            warning.level = warning.level + 1
            warning.levelName = level_map.get(warning.level)
            is_change = True
        if is_change:
            try:
                task.save()
            except:
                print(task.id)


if __name__ == "__main__":
    update_task_level()
